import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

import input, preprocess

# 模型训练和预测
def train_and_predict(Features, year_to_predict, host_country, num_of_programs):
    # 特征选择
    features = ["Year", "Is_Host", "Last_Total_Medals", "Total_Programs", "Athlete_Win_Ratio"]
    
    # 奖牌总数
    X_total = Features[Features["Year"] < year_to_predict][features]
    y_total = Features[Features["Year"] < year_to_predict]["Total"]
    # 金牌数
    X_gold = Features[Features["Year"] < year_to_predict][features]
    y_gold = Features[Features["Year"] < year_to_predict]["Gold"]

    # 确保特征数据为数值类型
    X_total = X_total.apply(pd.to_numeric, errors='coerce').fillna(0)
    y_total = pd.to_numeric(y_total, errors='coerce').fillna(0)
    X_gold = X_gold.apply(pd.to_numeric, errors='coerce').fillna(0)
    y_gold = pd.to_numeric(y_gold, errors='coerce').fillna(0)

    # 数据分割
    year_separate = 2000
    X_train_total = X_total[X_total["Year"] <= year_separate]
    y_train_total = y_total[X_total["Year"] <= year_separate]
    X_val_total = X_total[X_total["Year"] > year_separate]
    y_val_total = y_total[X_total["Year"] > year_separate]

    X_train_gold = X_gold[X_gold["Year"] <= year_separate]
    y_train_gold = y_gold[X_gold["Year"] <= year_separate]
    X_val_gold = X_gold[X_gold["Year"] > year_separate]
    y_val_gold = y_gold[X_gold["Year"] > year_separate]

    # 模型训练（奖牌总数）
    model_total = RandomForestRegressor(n_estimators=100, random_state=42)
    model_total.fit(X_train_total, y_train_total)
    # 模型评估（奖牌总数）
    y_pred_total = model_total.predict(X_val_total)
    mse_total = mean_squared_error(y_val_total, y_pred_total)
    print(f"\nMean Squared Error (Total Medals Validation Set): \033[91m{mse_total}\033[0m")

    # 模型训练（金牌数）
    model_gold = RandomForestRegressor(n_estimators=100, random_state=42)
    model_gold.fit(X_train_gold, y_train_gold)
    # 模型评估（金牌数）
    y_pred_gold = model_gold.predict(X_val_gold)
    mse_gold = mean_squared_error(y_val_gold, y_pred_gold)
    print(f"\nMean Squared Error (Gold Medals Validation Set): \033[91m{mse_gold}\033[0m")

    # 获取本次奥运会的特征数据
    latest_data = Features[Features["Year"] == (year_to_predict - 4)].copy()
    latest_data["Year"] = year_to_predict
    latest_data["Is_Host"] = (latest_data["Country"] == host_country).astype(int)
    latest_data["Last_Total_Medals"] = latest_data["Total"]
    latest_data["Total_Programs"] = num_of_programs
    latest_data["Athlete_Win_Ratio"] = latest_data["Athlete_Win_Ratio"]

    # 用本次奥运会的特征数据作为x来预测成绩
    future_X = latest_data[features]
    predictions_total = model_total.predict(future_X)
    predictions_gold = model_gold.predict(future_X)

    # 得到奖牌总数和金牌数
    latest_data["Total_Medals"] = predictions_total.astype(float).round(2)
    latest_data["Gold_Medals"] = predictions_gold.astype(float).round(2)

    return latest_data


# 评价国家实力
def evaluate_strength(latest_data, features, year_to_predict):
    last_olympics = features[features["Year"] == (year_to_predict - 4)].copy()
    # 使用熵权法计算国家实力
    data_for_entropy = last_olympics[["Total", "Gold", "Total_Programs", "Last_Total_Medals", "Is_Host"]]
    # 标准化数据
    normalized_data = (data_for_entropy - data_for_entropy.min()) / (data_for_entropy.max() - data_for_entropy.min())
    p = normalized_data / normalized_data.sum()
    entropy = (-p * np.log(p)).sum() # 计算熵值
    p = p.replace(0, 1e-10) # 替换零值
    weights = (1 - entropy) / (1 - entropy).sum() # 计算权重

    # 计算Strength
    last_olympics["Strength"] = (last_olympics["Total"] * weights["Total"] + 
                                 last_olympics["Gold"] * weights["Gold"] +
                                 last_olympics["Total_Programs"] * weights["Total_Programs"] +
                                 last_olympics["Last_Total_Medals"] * weights["Last_Total_Medals"] +
                                 last_olympics["Is_Host"] * weights["Is_Host"]).astype(float).round(1)
    
    latest_data["Strength"] = (latest_data["Total_Medals"] * weights["Total"] + 
                               latest_data["Gold_Medals"] * weights["Gold"] +
                               latest_data["Total_Programs"] * weights["Total_Programs"] +
                               latest_data["Last_Total_Medals"] * weights["Last_Total_Medals"] +
                               latest_data["Is_Host"] * weights["Is_Host"]).astype(float).round(1)
    
    # 合并预测数据和上一次奥运会的数据
    merged_data = latest_data.merge(last_olympics[["Country", "Strength"]], on="Country", suffixes=("", "_Last"))
    # 计算进步或退步
    merged_data["Progress"] = (merged_data["Strength"] - merged_data["Strength_Last"]).astype(float).round(1)

    # 按预测奖牌总数降序排序
    return merged_data.sort_values(by="Total_Medals", ascending=False)[["Country", "Total_Medals", "Gold_Medals", "Strength", "Strength_Last", "Progress"]]


# 主程序
if __name__ == "__main__":
    file_path = "./data/"  # 替换为你的文件路径
    dataframes = input.load_data(file_path)

    features = preprocess.preprocess_data(dataframes)
    year_to_predict = 2028
    host_country = "United States"
    num_of_programs = 400

    # 预测
    predict_data = train_and_predict(features, year_to_predict, host_country, num_of_programs)

    # 评价
    predictions = evaluate_strength(predict_data, features, year_to_predict)

    # 打印输出
    print("\nPredicted Medals for", year_to_predict, "Olympics:")
    print(predictions)

    predictions.to_csv("./predictions.csv", index=False, encoding="utf-8")
    print("预测结果已保存到 predictions.csv 文件中。")